import time
import zmq
import config
import MDP
from utils import debug_log

class MajorWorker(object):
    HEARTBEAT_LIVENESS = config.HEARTBEAT_LIVENESS
    broker = None
    ctx = None
    service = None
    worker = None       # Socket to broker

    heartbeat_at = 0    # When to send HEARTBEAT (relative to time.time(), so in seconds)
    liveness = 0        # How many attempts left
    heartbeat = config.HEARTBEAT    # Heartbeat delay, msecs
    reconnect = config.RECONNECT    # Reconnect delay, msecs
    timeout = config.TIMEOUT        # poller timeout

    # Return address, if any
    reply_to = None

    def __init__(self, broker, service):
        self.broker = broker
        self.service = service
        self.ctx = zmq.Context()
        self.poller = zmq.Poller()
        self.reconnect_to_broker()

    def reconnect_to_broker(self):
        """Connect or reconnect to broker"""
        if self.worker:
            self.poller.unregister(self.worker)
            self.worker.close()

        self.worker = self.ctx.socket(zmq.DEALER)
        self.worker.linger = 0
        self.worker.connect(self.broker)
        self.poller.register(self.worker, zmq.POLLIN)

        # Register service with broker
        self.send_to_broker(MDP.W_READY, self.service, [])

        # If liveness hits zero, queue is considered disconnected
        self.liveness = self.HEARTBEAT_LIVENESS
        self.heartbeat_at = time.time() + 1e-3 * self.heartbeat

    def send_to_broker(self, command, option=None, msg=None):
        """Send message to broker.
        If no msg is provided, creates one internally
        """
        if msg is None:
            msg = []
        elif not isinstance(msg, list):
            msg = [msg]

        if option:
            msg = [option] + msg

        msg = ['', MDP.W_WORKER, command] + msg
        print "send msg: ", msg
        self.worker.send_multipart(msg)

    def recv(self):
        """Send reply, if any, to broker and wait for next request."""

        # Poll socket for a reply, with timeout
        while True:
            items = self.poller.poll(self.timeout)
            if items:
                msg = self.worker.recv_multipart()
                debug_log("recv: ", msg)

                self.liveness = self.HEARTBEAT_LIVENESS

                if len(msg) < 3:
                    # TODO
                    break;

                empty = msg.pop(0)
                if empty != '':
                    # TODO
                    break;

                header = msg.pop(0)
                if header != MDP.W_WORKER:
                    # TODO
                    break;

                command = msg.pop(0)
                if command == MDP.W_REQUEST:
                    # We should pop and save as many addresses as there are
                    # up to a null part, but for now, just save one
                    self.reply_to = msg.pop(0)

                    # pop empty
                    if msg.pop(0) != '':
                        # TODO
                        break;

                    return msg # We have a request to process
                elif command == MDP.W_HEARTBEAT:
                    # Do nothing for heartbeats
                    pass
                elif command == MDP.W_DISCONNECT:
                    self.reconnect_to_broker()
                else :
                    pass
                    # logging.error("E: invalid input message: ")
            else:
                self.liveness -= 1
                if self.liveness == 0:
                    # logging.warn("W: disconnected from broker - retrying")
                    time.sleep(1e-3*self.reconnect)
                    self.reconnect_to_broker()

            # Send HEARTBEAT if it's time
            if time.time() > self.heartbeat_at:
                self.send_to_broker(MDP.W_HEARTBEAT)
                self.heartbeat_at = time.time() + 1e-3*self.heartbeat
        return None

    def destroy(self):
        self.ctx.destroy(0)

